/*
    Author: Christopher Diggins
    License: MIT Licence 1.0
    Website: http://www.heron-language.com
*/

#ifndef LLSCHEME_NAME_REMOVER_HPP
#define LLSCHEME_NAME_REMOVER_HPP

namespace lls
{
	struct NameRemover 
	{
		Strings vals;
		Strings defs;
		int depth;
		bool verbose;

		NameRemover(bool bVerbose = true) 
			: depth(0), verbose(bVerbose)
		{		
			if (verbose)
				printf("\nName Removal\n");
		}
		
		void PushTmp() {
			vals.Push(GenSym());
		}

		int Lookup(ConstString var) {
			for (int i=vals.Count() - 1; i >= 0; --i) {
				if (var == vals[i])
					return vals.Count() - 1 - i;
			}
			return -1;
		}

		int LookupDef(ConstString var) {
			for (int i=defs.Count() - 1; i >= 0; --i) {
				if (var == defs[i])
					return i;
			}
			return -1;
		}

		void RemoveVar(int n) {
			vals.Remove(vals.Count() - 1 - n);
		}

		void RemoveNamesFromChildren(SExpr* src, SExpr* dest) {
			for (int i=0; i < src->GetNumChildren(); ++i) {
				SExpr* oldChild = src->GetChild(i);
				SExpr* newChild = (*this)(oldChild);
				dest->AddChild(newChild);
			}
		}

		bool IsUnnamedVar(ConstString s) {
			return s.size() > 0 && s[0] == '$';
		}

		SExpr* RemoveNames(SExpr* x) {
			int valcnt = vals.Count();
			if (x->Is<Module>()) {
				Module* mod = x->As<Module>();
				Assert(false, "not implemented");
				// TODO: handle classes properly.
				/*
				Seq* seq = mod->GetStatements();
				// Scan the names of all the defines. 
				for (int i=0; i < seq->GetNumChildren(); ++i) {
					SExpr* child = seq->GetChild(i);					
					if (child->Is<Define>()) {
						Define* def = child->As<Define>();
						defs.Push(def->GetNameStr());
					}
					else {
						Error("Modules can only have functions as children");
					}
				}
				*/
				Module* r = new Module();				
				RemoveNamesFromChildren(x, r);
				return r;
				
			}
			else if (x->Is<NoOp>()) {
				return x->Clone();
			}
			else if (x->Is<Name>()) {
				return x->Clone();
			}
			else if (x->Is<Type>()) {
				return x->Clone();
			}
			else if (x->Is<Define>()) {
				Define* r = new Define();
				RemoveNamesFromChildren(x, r);
				Assert(vals.Count() == valcnt);
				return r;
			}			
			else if (x->Is<Formals>()) {
				Formals* r = new Formals();
				RemoveNamesFromChildren(x, r);
				return r;
			}
			else if (x->Is<Formal>()) {
				Formal* formal = x->As<Formal>();
				vals.Push(formal->GetNameStr());
				return formal->Clone();
			}
			else if (x->Is<Args>()) {
				Args* r = new Args();
				RemoveNamesFromChildren(x, r);
				return r;
			}
			else if (x->Is<Apply>()) {
				Apply* r = new Apply();
				RemoveNamesFromChildren(x, r);				
				Args* args = r->As<Apply>()->GetArgs();
				assert(vals.Count() == valcnt + args->GetNumChildren() + 1);
				for (int i=0; i < args->GetNumChildren(); ++i) {
					vals.Pop();
				}
				assert(vals.Count() == valcnt + 1);
				vals.Pop(); // the function pointer 
				vals.Push("$result$");
				return r;
			}
			else if (x->Is<Invoke>()) {
				Invoke* r = new Invoke();
				RemoveNamesFromChildren(x, r);
				Args* args = r->GetArgs(); 
				Assert(vals.Count() == valcnt + args->GetNumChildren() + 1);
				for (int i=0; i < args->GetNumChildren() + 1; ++i) {
					vals.Pop();
				}
				Assert(vals.Count() == valcnt);
				vals.Push("$result$");
				Assert(vals.Count() == valcnt + 1);
				return r;
			}
			else if (x->Is<If>()) {
				If* r = new If();	
				Assert(x->GetNumChildren() == 3);
				r->AddChild((*this)(x->GetChild(0)));
				Assert(vals.Count() == valcnt + 1);
				vals.Pop();				
				r->AddChild((*this)(x->GetChild(1)));
				if (vals.Count() == valcnt + 1) 
					vals.Pop(); 
				else 
					Assert(vals.Count() == valcnt);
				r->AddChild((*this)(x->GetChild(2)));
				if (vals.Count() == valcnt + 1) 
					vals.Pop(); 
				else 
					Assert(vals.Count() == valcnt);
				return r;
			}
			else if (x->Is<While>()) {
				While* r = new While();				
				Assert(x->GetNumChildren() == 2);
				// Condition
				r->AddChild((*this)(x->GetChild(0)));
				vals.Pop();				
				r->AddChild((*this)(x->GetChild(1)));
				return r;
			}
			else if (x->Is<Lambda>()) {
				Error("todo");
				return NULL;
			}
			else if (x->Is<Return>()) {
				Return* r = new Return();
				int n = vals.Count();
				RemoveNamesFromChildren(x, r);
				Assert(vals.Count() == n + 1);
				vals.Pop();
				return r;
			}
			else if (x->Is<For>()) {
				For* f = x->As<For>();
				For* r = new For();
				String s = f->GetNameStr();
				vals.Push(s);
				r->AddChild((*this)(f->GetNameExpr()));
				r->AddChild((*this)(f->GetType()));
				int n = vals.Count();
				r->AddChild((*this)(f->Init()));
				Assert(vals.Count() == n + 1);
				vals.Pop();
				r->AddChild((*this)(f->Inv()));
				Assert(vals.Count() == n + 1);
				vals.Pop();
				r->AddChild((*this)(f->Next()));
				Assert(vals.Count() == n);
				r->AddChild((*this)(f->Body()));
				vals.Pop(); 
				Assert(vals.Count() == valcnt);
				return r;
			}
			else if (x->Is<ForEach>()) {
				ForEach* f = x->As<ForEach>();
				ForEach* r = new ForEach();
				r->AddChild((*this)(f->GetNameExpr()));
				r->AddChild((*this)(f->GetType()));
				int n = vals.Count();
				r->AddChild((*this)(f->List()));
				vals.Pop();
				String s = f->GetNameStr();
				vals.Push(s);
				r->AddChild((*this)(f->Body()));
				vals.Pop();
				return r;
			}
			else if (x->Is<Final>()) {
				String sVar = x->As<Final>()->GetNameStr();
				int n = Lookup(sVar);
				if (n < 0)
					Error("Could not find variable " + sVar);
				RemoveVar(n);
				return new DigN(n, sVar);
			}
			else if (x->Is<Undeclare>()) {
				String sVar = x->As<Undeclare>()->GetNameStr();
				int n = Lookup(sVar);
				if (n < 0)
					Error("Could not find variable " + sVar);
				RemoveVar(n);
				return new ZapN(n, sVar);				
			}
			else if (x->Is<Var>()) {
				SExpr* r = (*this)(x->As<Var>()->GetVal());
				Assert(vals.Count() == valcnt + 1);
				vals.Pop();
				String var = x->As<Var>()->GetNameStr();
				vals.Push(var);
				return r;
			}
			else if (x->Is<GetField>()) {
				GetField* t = x->As<GetField>();
				GetField* r = new GetField();
				r->AddChild((*this)(t->Object()));				
				Assert(vals.Count() == valcnt + 1);
				vals.Pop();
				String s = t->GetNameStr();
				r->AddChild(t->GetNameExpr()->Clone());
				Assert(vals.Count() == valcnt);
				PushTmp();
				Assert(vals.Count() == valcnt + 1);
				return r;
			}
			else if (x->Is<SetField>()) {
				SetField* t = x->As<SetField>();
				SetField* r = new SetField();
				r->AddChild((*this)(t->Object()));				
				Assert(vals.Count() == valcnt + 1);
				vals.Pop();
				String s = t->GetNameStr();
				r->AddChild(t->GetNameExpr()->Clone());
				Assert(vals.Count() == valcnt);
				PushTmp();
				r->AddChild(t->GetVal());
				vals.Pop();
				Assert(vals.Count() == valcnt);
				return r;
			}
			else if (x->Is<Pop>()) {
				vals.Pop();
				return x->Clone();
			}
			else if (x->Is<Get>()) {
				String var = x->As<Get>()->GetNameStr();
				int n = Lookup(var);
				if (n >= 0) {
					GetN* r = new GetN(n, var);
					vals.Push("$temporary$");
					return r;
				}
				n = LookupDef(var);
				if (n >= 0) {
					GetF* r = new GetF(n, var);
					vals.Push("$definition$");
					return r;
				}		
				else {
					vals.Push("$external$");
					return new External(var);
				}
				return NULL;
			}
			else if (x->Is<Set>()) {
				String var = x->As<Set>()->GetNameStr();
				(*this)(x->As<Set>()->GetVal());
				int n = Lookup(var);
				if (n < 0) {
					Error("Could not find name " + var);
				}
				SetN* r = new SetN(n, var);
				vals.Pop();
				return r;
			}
			else if (x->Is<Primitive>()) {
				return ProcessPrimitive(x->As<Primitive>());
			}
			else if (x->Is<Literal>()) {
				Literal* lit = x->As<Literal>();
				String val = lit->Label();
				vals.Push(val);
				return lit->Clone();
			}
			else if (x->Is<Seq>()) {
				Seq* r = new Seq();
				for (int i=0; i < x->GetNumChildren(); ++i) {
					SExpr* child = x->GetChild(i);
					if (child->Is<Return>()) {
						r->AddChild((*this)(child));
					}
					else if (child->Is<Var>()) {
						r->AddChild((*this)(child));
						Assert(vals.Count() == valcnt + 1);
						++valcnt;
					}
					else if (child->Is<Undeclare>()) {
						r->AddChild((*this)(child));
						--valcnt;
					}
					else {
						// All other statement types should leave 
						// nothing on the stack, so we add extra Pop
						// statements as neccessary
						int n = vals.Count();
						r->AddChild((*this)(child));
						if (vals.Count() > n) {
							Assert(vals.Count() == n + 1);
							r->AddChild(new Pop());
							String s = vals.Top();
							Assert(IsUnnamedVar(s));
							vals.Pop();
						}
					}	
					Assert(vals.Count() == valcnt);
				}
				return r;
			}
			else if (x->Is<GetAt>()) {
				GetAt* r = new GetAt();
				RemoveNamesFromChildren(x, r);
				vals.Pop();
				return r;
			}
			else if (x->Is<SetAt>()) {
				SetAt* r = new SetAt();
				RemoveNamesFromChildren(x, r);
				vals.Pop();
				vals.Pop();
				vals.Pop();
				return r;
			}
			else if (x->Is<New>()) {
				New* r = new New();
				RemoveNamesFromChildren(x, r);
				while (vals.Count() > valcnt) {
					vals.Pop();
				}
				vals.Push("$new$");
				Assert(vals.Count() == valcnt + 1);
				return r;
			}
			else {
				Error(String("Can not remove names from the s-expression ") + x->ToString());
			}

			Error("missing return statement");
			return NULL;
		}

		SExpr* operator()(SExpr* x) {
			if (verbose) {
				for (int i=0; i < depth; ++i)
					printf("  ");
				if (vals.Count() > 0) {
					printf("%d %s %s\n", vals.Count(), vals.Top().c_str(), x->ToString().c_str());
				}
				else {
					printf("0 _empty_ %s\n", x->Label());
				}
			}
			depth += 1;
			SExpr* r = RemoveNames(x);
			depth -= 1;
			return r;
		}

		SExpr* PrimChecker(Primitive* x, ConstString out, ConstString a0) {
			Primitive* r = new Primitive(x->id);
			RemoveNamesFromChildren(x, r);
			return r;
		}

		SExpr* PrimChecker(Primitive* x, ConstString out, ConstString a0, ConstString a1) {
			Primitive* r = new Primitive(x->id);
			RemoveNamesFromChildren(x, r);
			vals.Pop();
			return r;
		}

		SExpr* PrimChecker(Primitive* x, ConstString out, ConstString a0, ConstString a1, ConstString a2) {
			Primitive* r = new Primitive(x->id);
			RemoveNamesFromChildren(x, r);
			vals.Pop();
			vals.Pop();
			return r;
		}

		SExpr* ProcessPrimitive(Primitive* prim) {
			switch (prim->Id()) 
			{
			case prim_inc: return PrimChecker(prim, "num", "num");
			case prim_dec: return PrimChecker(prim, "num", "num");
			case prim_neg: return PrimChecker(prim, "num", "num");
			case prim_add: return PrimChecker(prim, "num", "num", "num");
			case prim_sub: return PrimChecker(prim, "num", "num", "num");
			case prim_mul: return PrimChecker(prim, "num", "num", "num");
			case prim_div: return PrimChecker(prim, "num", "num", "num");
			case prim_mod: return PrimChecker(prim, "num", "num", "num");
			case prim_gt: return PrimChecker(prim, "num", "num", "bool");
			case prim_lt: return PrimChecker(prim, "num", "num", "bool");
			case prim_gteq: return PrimChecker(prim, "num", "num", "bool");
			case prim_lteq: return PrimChecker(prim, "num", "num", "bool");
			case prim_eq: return PrimChecker(prim, "any", "any", "bool");
			case prim_neq: return PrimChecker(prim, "any", "any", "bool");
			case prim_and: return PrimChecker(prim, "bool", "bool", "bool");
			case prim_or: return PrimChecker(prim, "bool", "bool", "bool");
			case prim_xor: return PrimChecker(prim, "bool", "bool", "bool");
			case prim_cond: return PrimChecker(prim, "any", "bool", "any", "any");
			case prim_null: vals.Push("$null$"); return prim->Clone();
			case prim_nil: vals.Push("$nil$"); return prim->Clone();
			case prim_unused: Error("unused primitive"); return NULL;
			default: Error("Unrecognized enumeration " + ToStr(prim->Id())); return NULL;
			}
		}
	};

	bool HasNames(SExpr* x) {
		if (x->Is<Get>()) {
			return true;
		}
		else if (x->Is<Set>()) {
			return true;
		}
		else {
			return x->QueryAny(HasNames);
		}
	}

	Module* RemoveNames(Module* x, bool verbose) {
		NameRemover nr(verbose);
		return dynamic_cast<Module*>(nr(x));
	}
}

#endif
